package com.pro01.apirednote.service.serviceImpl;

import com.hankcs.hanlp.HanLP;
import com.hankcs.hanlp.seg.common.Term;
import com.pro01.apirednote.service.ChineseVectorService;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.word2vec.Word2Vec;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.ops.transforms.Transforms;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
@Service
public class ChineseVectorServiceImpl implements ChineseVectorService {

    @Value("${vector.model.path:}")
    private String modelPath;

    private Word2Vec word2Vec;
    private final AtomicBoolean modelLoaded = new AtomicBoolean(false);
    private final int vectorSize = 300; // 词向量维度

    @PostConstruct
    public void init() {
        // 如果配置了模型路径，则尝试加载模型
        if (modelPath != null && !modelPath.isEmpty()) {
            try {
                initWordVectorModel();
            } catch (Exception e) {
                log.error("初始化词向量模型失败", e);
            }
        }
    }

    @Override
    public void initWordVectorModel() {
        if (modelPath == null || modelPath.isEmpty()) {
            log.warn("词向量模型路径未配置，无法加载模型");
            return;
        }

        try {
            File modelFile = new File(modelPath);
            if (!modelFile.exists()) {
                log.error("词向量模型文件不存在: {}", modelPath);
                return;
            }

            log.info("开始加载词向量模型: {}", modelPath);
            word2Vec = WordVectorSerializer.readWord2VecModel(modelFile);
            modelLoaded.set(true);
            log.info("词向量模型加载完成，词汇量: {}", word2Vec.vocab().numWords());
        } catch (Exception e) {
            log.error("加载词向量模型失败", e);
            modelLoaded.set(false);
        }
    }

    @Override
    public double[] textToVector(String text) {
        if (!modelLoaded.get() || text == null || text.isEmpty()) {
            return new double[vectorSize]; // 返回零向量
        }

        try {
            // 使用HanLP进行分词
            List<Term> terms = HanLP.segment(text);
            List<String> words = new ArrayList<>();
            for (Term term : terms) {
                words.add(term.word);
            }

            // 计算文本向量（词向量的平均）
            INDArray textVector = Nd4j.zeros(vectorSize);
            int wordCount = 0;

            for (String word : words) {
                if (word2Vec.hasWord(word)) {
                    INDArray wordVector = word2Vec.getWordVectorMatrix(word);
                    textVector.addi(wordVector);
                    wordCount++;
                }
            }

            // 如果文本中有词在词向量模型中，则取平均
            if (wordCount > 0) {
                textVector.divi(wordCount);
            }

            // 转换为普通数组
            double[] result = new double[vectorSize];
            for (int i = 0; i < vectorSize; i++) {
                result[i] = textVector.getDouble(i);
            }
            return result;
        } catch (Exception e) {
            log.error("文本向量化失败: {}", e.getMessage());
            return new double[vectorSize]; // 返回零向量
        }
    }

    @Override
    public double calculateSimilarity(String text1, String text2) {
        if (!modelLoaded.get()) {
            log.warn("词向量模型未加载，无法计算相似度");
            return 0.0;
        }

        try {
            // 获取两段文本的向量表示
            double[] vector1 = textToVector(text1);
            double[] vector2 = textToVector(text2);

            // 转换为INDArray计算余弦相似度
            INDArray ind1 = Nd4j.create(vector1);
            INDArray ind2 = Nd4j.create(vector2);

            // 计算余弦相似度
            double similarity = Transforms.cosineSim(ind1, ind2);

            // 确保结果在0-1之间
            return Math.max(0, Math.min(1, similarity));
        } catch (Exception e) {
            log.error("计算文本相似度失败: {}", e.getMessage());
            return 0.0;
        }
    }

    @Override
    public boolean isModelLoaded() {
        return modelLoaded.get();
    }
} 